import math
import os

import torch

from config import *
import numpy as np
from torch.utils.data import Dataset
from PIL import Image
from utils import *
from torchvision import transforms

tf=transforms.Compose([
    transforms.ToTensor()
])

def one_hot(cls_num,i):
    # rst = np.zeros(cls_num) 用于创建一个由 0 组成的 numpy 数组 rst，其长度为 cls_num。
    rst=np.zeros(cls_num)
    rst[i]=1
    return rst

class YoloDataSet(Dataset):
    def __init__(self):
        f=open('data.txt','r')
        self.dataset=f.readlines()  # 按行读取文件的内容。返回结果是一个包含每行内容的列表。

    def __len__(self):
        return len(self.dataset)

    def __getitem__(self, index):
        data=self.dataset[index].strip()   # data为单张图片对应的数据
        temp_data=data.split()             # 根据空格拆分字符串为列表
        # 这 temp_data 的列表中的元素转换为浮点型，并将结果存储在名为 _boxes 的 NumPy 数组中。
        # 在转换时，代码示例使用了列表切片 temp_data[1:] 来跳过列表中的第一个元素。
        _boxes=np.array([float(x) for x in temp_data[1:]])
        # 使用 np.split() 函数将名为 _boxes 的 NumPy 数组拆分成多个子数组，并将结果存储在名为 boxes 的列表中。除以5是因为{cls cx cy w h}
        # boxes为存储多个数组变量的列表变量。其中的每个数组变量为一张图片中一个object对应的标签
        boxes=np.split(_boxes,len(_boxes)//5)
        img=make_416_image(os.path.join('data\\images',temp_data[0]))
        w,h=img.size
        case=416/w         # case为需要缩放的比例，如果图片过大就需要缩放
        img=img.resize((DATA_WIDTH,DATA_HEIGHT))
        img_data=tf(img)   # 将PIL图像对象转换为PyTorch图像张量
        labels={}          # 这行代码创建了一个空的字典对象
        # feature_size取值 13 26 52
        # _anchors取值 大框 [[168,302], [57,221], [336,284]] 中框 [[175,225], [279,160], [249,271]] 小框 [[129,209], [85,413], [44,42]]
        for feature_size,_anchors in ANCHORS.items():
            # 13 : 13*13*3*(5+CLASS_NUM)
            # 26 : 26*26*3*(5+CLASS_NUM)
            # 52 : 52*52*3*(5+CLASS_NUM)
            labels[feature_size]=np.zeros(shape=(feature_size,feature_size,3,5+CLASS_NUM))

            for box in boxes:     # 对于一张图片中的所有object的ground truth进行处理
                cls,cx,cy,w,h=box
                # 对于原来框的 cx cy w h 成比例缩放，使得该框在缩放后的图像上依然能框住该object
                cx,cy,w,h=cx*case,cy*case,w*case,h*case
                # modf用于将一个浮点数拆分为其小数部分和整数部分，_x为小数部分，x_index是其整数部分。
                # cx*feature_size/DATA_WIDTH  = cx / (DATA_WIDTH / feature_size)
                # (x_index,y_index)为该object对应的ground truth的中心点所在的grid cell的左上角坐标值。相对于整张图片，左上角(0,0)来计算的
                # (_x,_y)为该ground truth的中心点相对于所在grid cell的左上角坐标的偏移量
                _x,x_index=math.modf(cx*feature_size/DATA_WIDTH)
                _y,y_index=math.modf(cy*feature_size/DATA_HEIGHT)
                # 使用了 enumerate() 函数来迭代一个名为 _anchors 的列表并获取每个元素和其对应的索引
                # 如果feature_size大小为13，则_anchors值为[[168,302], [57,221], [336,284]]
                # i = 0, anchor = [168,302]
                # i = 1, anchor = [57,221]
                # i = 2, anchor = [336,284]
                for i,anchor in enumerate(_anchors):
                    area=w*h
                    # 这里近似地计算，真实框和Anchor的iou值
                    iou=min(area,ANCHORS_AREA[feature_size][i])/max(area,ANCHORS_AREA[feature_size][i])
                    # 计算出真实框的 w h 相比较 Anchor的 w h 的偏差比例
                    p_w, p_h = w / anchor[0], h / anchor[1]
                    # 这段代码是用于标记目标检测中的ground truth的标签
                    # np.array([]) 构成了一个大小为(5+CLASS_NUM)的数组，表示当前ground truth的标签信息。
                    # 该数组被赋值给 labels[feature_size][int(x_index),int(y_index),i]
                    # labels[feature_size]=np.zeros(shape=(feature_size,feature_size,3,5+CLASS_NUM))
                    # int(x_index)和int(y_index)是grid cell的索引值，它们代表了该grid cell在二维平面上的位置。
                    # 其中 i 是grid cell处Anchor的索引，一个grid cell生成三个Anchor
                    # 当 w = anchor[0] h = anchor[1]时，p_w p_h的值为1，进行对数变换后，其值为0。这个值越接近0越好，因为后面我们要计算损失。
                    labels[feature_size][int(x_index),int(y_index),i]=np.array([1.0,_x,_y,np.log(p_w),np.log(p_h),*one_hot(CLASS_NUM,int(cls))])

        return labels[13],labels[26],labels[52],img_data

if __name__ == '__main__':
    dataset=YoloDataSet()
    print(dataset[0][3].shape)    # ([3, 416, 416])
    print(dataset[0][0].shape)    # (13, 13, 3, 8)
    print(dataset[0][1].shape)    # (26, 26, 3, 8)
    print(dataset[0][2].shape)    # (52, 52, 3, 8)
